import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation
import heapq
import random
from datetime import datetime
import matplotlib

# ==================== 中文显示配置 ====================
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文显示
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# ==================== 系统配置 ====================
FLOORS = 8                    # 总楼层数
GRID_SIZE = (15, 15)          # 每层网格尺寸
EXIT_POS = (0, 14, 14)        # 出口位置（底层右下角）
PEOPLE_PER_FLOOR = 20         # 每层人数
BASE_SPEED = 0.8              # 基础移动速度
FRAME_DELAY = 100             # 动画帧间隔(ms)
SIMULATION_TIME = 500         # 最大模拟时间步

# ==================== 建筑结构生成 ====================
def create_building():
    """生成完全连通的三维建筑结构"""
    building = np.zeros((FLOORS, GRID_SIZE[0], GRID_SIZE[1]), dtype=int)
    
    # 创建双向通道系统
    for z in range(FLOORS):
        # 横向主通道（第7列）
        building[z, 7, :] = 1
        # 纵向主通道（第7行）
        building[z, :, 7] = 1
        
        # 双侧楼梯（第0-2列和第12-14列）
        building[z, 0:3, :] = 2
        building[z, 12:15, :] = 2
        
        # 连接楼梯和通道
        building[z, 3:6, 7] = 1  # 左侧通道连接
        building[z, 9:12, 7] = 1  # 右侧通道连接
        
    return building

# ==================== 路径规划系统 ====================
class Navigator:
    def __init__(self, building):
        self.building = building
        self.cost_field = np.full(building.shape, np.inf)
        self.calculate_cost_field()
        
    def calculate_cost_field(self):
        """使用优化的Dijkstra算法计算三维成本场"""
        heap = []
        heapq.heappush(heap, (0, EXIT_POS))
        self.cost_field[EXIT_POS] = 0
        
        while heap:
            cost, (z, x, y) = heapq.heappop(heap)
            
            # 26邻域搜索
            for dz in [-1, 0, 1]:
                for dx in [-1, 0, 1]:
                    for dy in [-1, 0, 1]:
                        if dz==dx==dy==0: continue
                        
                        nz, nx, ny = z+dz, x+dx, y+dy
                        if self.is_valid(nz, nx, ny) and self.is_connected(z, x, y, nz, nx, ny):
                            move_cost = self.get_move_cost(z, x, y, nz, nx, ny)
                            new_cost = cost + move_cost
                            
                            if new_cost < self.cost_field[nz, nx, ny]:
                                self.cost_field[nz, nx, ny] = new_cost
                                heapq.heappush(heap, (new_cost, (nz, nx, ny)))
                                
    def is_valid(self, z, x, y):
        """坐标有效性验证"""
        return (0 <= z < FLOORS and 
                0 <= x < GRID_SIZE[0] and 
                0 <= y < GRID_SIZE[1])
    
    def is_connected(self, oz, ox, oy, nz, nx, ny):
        """增强的连通性验证"""
        # 允许在楼梯区域跨层
        if oz != nz:
            return (self.building[oz, ox, oy] == 2 
                    and self.building[nz, nx, ny] == 2
                    and ox//12 == nx//12)  # 确保同侧楼梯
        
        # 本层移动需要通道或楼梯
        return self.building[oz, ox, oy] in [1, 2]
    
    def get_move_cost(self, oz, ox, oy, nz, nx, ny):
        """精确的移动成本计算"""
        # 楼梯移动成本
        if oz != nz:
            return 4.0  # 楼梯层间移动成本
        
        # 通道移动成本
        if self.building[oz, ox, oy] == 1:
            return 1.0  # 主通道
        else:
            return 1.5  # 连接通道

# ==================== 人员实体 ====================
class Person:
    def __init__(self, start_z, building):
        self.z = start_z
        self.building = building
        self.x, self.y = self.find_valid_position()
        self.speed = BASE_SPEED * random.uniform(0.7, 1.3)
        self.arrived = False
        self.path = []
        
    def find_valid_position(self):
        """确保生成在可达位置"""
        while True:
            x = random.randint(3, GRID_SIZE[0]-4)
            y = random.randint(3, GRID_SIZE[1]-4)
            if self.building[self.z, x, y] == 1:
                return x, y
            
    def move_step(self, nav):
        """智能路径跟踪"""
        if self.arrived: return
        
        current_cost = nav.cost_field[self.z, self.x, self.y]
        
        # 寻找最优移动方向
        best_move = None
        min_cost = current_cost
        
        for dz in [-1, 0, 1]:
            for dx in [-1, 0, 1]:
                for dy in [-1, 0, 1]:
                    if dz==dx==dy==0: continue
                    nz, nx, ny = self.z+dz, self.x+dx, self.y+dy
                    if nav.is_valid(nz, nx, ny) and nav.is_connected(*self.pos, nz, nx, ny):
                        if nav.cost_field[nz, nx, ny] < min_cost:
                            min_cost = nav.cost_field[nz, nx, ny]
                            best_move = (dz, dx, dy)
        
        if best_move:
            dz, dx, dy = best_move
            self.z += dz
            self.x += dx
            self.y += dy
            self.path.append((self.z, self.x, self.y))
            
            # 应用体力消耗
            self.apply_fatigue()
            
        # 到达检测
        if (self.z, self.x, self.y) == EXIT_POS:
            self.arrived = True
    
    def apply_fatigue(self):
        """精确的体力消耗模型"""
        a, b, c = 0.71, 0.51, 1.0
        total_steps = len(self.path)
        total_dist = total_steps * 0.6
        total_time = total_steps / self.speed if self.speed > 0 else 0
        
        if total_time > 0:
            exponent = (total_dist**a) / (total_time**b)
            self.speed *= (1 - 0.01*c*np.exp(exponent) + 0.01)
            self.speed = max(self.speed, 0.4)  # 设置最低速度阈值

    @property
    def pos(self):
        return (self.z, self.x, self.y)

# ==================== 可视化系统 ====================
class Visual3D:
    def __init__(self, building):
        self.fig = plt.figure(figsize=(16, 12))
        self.ax = self.fig.add_subplot(111, projection='3d')
        self.building = building
        self.start_time = datetime.now()
        
    def draw_infrastructure(self):
        """三维基础设施绘制"""
        # 绘制各层平面
        for z in range(FLOORS):
            X, Y = np.meshgrid(range(GRID_SIZE[0]), range(GRID_SIZE[1]))
            Z = np.full_like(X, z)
            
            # 主通道高亮
            mask = (self.building[z] == 1)
            self.ax.scatter(X[mask], Y[mask], Z[mask], 
                           c='yellow', s=1, alpha=0.2)
            
            # 楼梯区域标注
            stairs_mask = (self.building[z] == 2)
            self.ax.scatter(X[stairs_mask], Y[stairs_mask], Z[stairs_mask],
                           c='purple', s=2, alpha=0.3)
            
        # 出口标注
        self.ax.scatter(EXIT_POS[2], EXIT_POS[1], EXIT_POS[0],
                       c='lime', s=200, marker='*', edgecolors='black')
        
    def update_display(self, people, frame):
        """实时更新显示"""
        self.ax.clear()
        self.draw_infrastructure()
        
        # 统计信息
        remaining = sum(not p.arrived for p in people)
        evacuated = len(people) - remaining
        
        # 绘制未撤离人员
        x = [p.x for p in people if not p.arrived]
        y = [p.y for p in people if not p.arrived]
        z = [p.z for p in people if not p.arrived]
        self.ax.scatter(x, y, z, c='red', s=20, alpha=0.8)
        
        # 绘制撤离轨迹
        for p in people:
            if len(p.path) > 1:
                path = np.array(p.path)
                self.ax.plot(path[:,2], path[:,1], path[:,0], 
                            c='blue', alpha=0.2, lw=0.5)
        
        # 显示统计信息
        stats_text = (f"总人数: {len(people)}\n"
                     f"已撤离: {evacuated}\n"
                     f"剩余: {remaining}\n"
                     f"运行时间: {(datetime.now()-self.start_time).total_seconds():.1f}秒")
        self.ax.text2D(0.02, 0.95, stats_text, transform=self.ax.transAxes,
                      bbox=dict(facecolor='white', alpha=0.9))
        
        # 视角控制
        self.ax.view_init(elev=35, azim=frame*1.5)
        self.ax.set_title(f"三维超高层建筑疏散模拟 - 时间步: {frame}")
        return self.ax,

# ==================== 初始化系统 ====================
print("正在初始化建筑结构...")
building = create_building()
print("正在计算导航场...")
nav = Navigator(building)
print("生成疏散人员...")
people = [Person(z, building) for z in range(1, FLOORS) for _ in range(PEOPLE_PER_FLOOR)]

# ==================== 运行动画 ====================
print("启动可视化引擎...")
vis = Visual3D(building)

def animate(frame):
    if frame >= SIMULATION_TIME:
        return vis.ax,
    
    print(f"\r正在计算: {frame+1}/{SIMULATION_TIME}", end='', flush=True)
    for p in people:
        if not p.arrived:
            p.move_step(nav)
    return vis.update_display(people, frame)

print("\n开始模拟...")
ani = FuncAnimation(vis.fig, animate, frames=SIMULATION_TIME, interval=FRAME_DELAY, blit=False)
plt.show()